import ssl
from threading import Lock
import websocket
from testcase.utils.Logger import Logger


class Wss:
    def __init__(self, wss_url):
        self.wss_url = wss_url
        self.log = Logger()
        self.last_message = None
        self.exit_code = None
        self._message_lock = Lock()
        self.execution_success = False
        self.ws = None

    def start_wss_connection(self):
        self.ws = websocket.WebSocketApp(
            self.wss_url,
            on_open=self.on_open,
            on_message=self.on_message,
            on_error=self.on_error,
            on_close=self.on_close
        )
        self.ws.run_forever(sslopt={"cert_reqs": ssl.CERT_NONE})

    def on_open(self, ws):
        self.log.info("已连接到 WSS 服务器")

    def on_message(self, ws, message):
        with self._message_lock:
            self.last_message = message
        self._process_message(message)

    def on_error(self, ws, error):
        self.log.info(error)

    def on_close(self, ws, close_status_code, close_msg):
        self.log.info(f'连接关闭:{close_status_code}')
        self.log.info(close_msg)

    def _process_message(self, message):
        """
        处理接收到的消息，提取退出码
        """
        if message.startswith("FINAL_EXECUTE_EXIT_CODE:"):
            try:
                self.exit_code = int(message.split(":")[1])
            except (IndexError, ValueError):
                self.log.error(f"无效的退出码格式: {message}")
                self.exit_code = None

    def assert_execution_success(self):
        """
        外部调用的断言方法
        """
        if self.exit_code is None:
            self.execution_success = False
            raise AssertionError("WSS 连接未完成或未收到退出码")
        if self.exit_code == 0:
            self.execution_success = True
        if self.exit_code == -1:
            self.execution_success = False
            raise AssertionError("命令执行失败")
        return self.execution_success

    def get_last_message(self):
        """
        获取最新的信息
        :return:
        """
        with self._message_lock:
            return self.last_message

    def stop_connection(self):
        """安全关闭连接"""
        if self.ws:
            self.ws.close()
